import numpy as np
import json
import os
from bishe_situations.utils import PYSR_PARAMS, GPLEARN_PARAMS

def generate_param_variants(base_params, n_variants=1000, variation_range=0.5):
    """
    生成参数变体
    base_params: 基础参数字典
    n_variants: 要生成的变体数量
    variation_range: 变化范围（上下浮动百分比）
    """
    variants = []
    
    for _ in range(n_variants):
        variant = {}
        for key, value in base_params.items():
            # 生成随机变化因子 (1 ± variation_range)
            factor = 1 + np.random.uniform(-variation_range, variation_range)
            # 应用变化
            new_value = value * factor
            # 对于概率类型的参数，确保在[0,1]范围内
            if 'probability' in key or 'p_' in key:
                new_value = np.clip(new_value, 0, 1)
            variant[key] = float(new_value)
        if base_params == GPLEARN_PARAMS:
            if np.sum(list(variant.values())) <= 1:
                variants.append(variant)
        else:
            variants.append(variant)
    return variants

def save_params_variants():
    # 创建保存目录
    os.makedirs('param_variants', exist_ok=True)
    
    # 生成PySR参数变体
    pysr_variants = generate_param_variants(PYSR_PARAMS)
    with open('param_variants/pysr_variants.json', 'w') as f:
        json.dump({
            'base_params': PYSR_PARAMS,
            'variants': pysr_variants
        }, f, indent=4)
    
    # 生成GPlearn参数变体
    gplearn_variants = generate_param_variants(GPLEARN_PARAMS)
    with open('param_variants/gplearn_variants.json', 'w') as f:
        json.dump({
            'base_params': GPLEARN_PARAMS,
            'variants': gplearn_variants
        }, f, indent=4)
    
    print(f"已生成 {len(pysr_variants)} 个PySR参数变体")
    print(f"已生成 {len(gplearn_variants)} 个GPlearn参数变体")
    print("参数变体已保存到 param_variants 目录")

if __name__ == "__main__":
    save_params_variants() 